from enum import unique

from myapp.const.base import TEnum


# Defines an enumeration for different kinds of tasks or operations.
@unique
class EnumKind(TEnum):
    """
    训练、推理资源池类型
    """

    # Represents a custom job task.
    job = 'job'  # 自定义任务
    # Represents an auto-learning task.
    auto_ml = 'auto_ml'  # 自动学习
    # Represents a finetuning task.
    finetune = 'finetune'  # 微调
    # Represents a post-pretraining task (incremental pre-training).
    post_pretraining = 'post_pretraining'  # 增量预训练
    # Represents a model compression task.
    compress = 'compress'  # 模型压缩
    # Represents a model merging task.
    merge = 'merge'  # 模型融合
    # Represents a Reinforcement Learning from Human Feedback (RLHF) task.
    rlhf = "rlhf" # 偏好对齐


# Checks if a given kind of task can result in the creation of a model.
# Args:
#   kind: The kind of task to check.
# Returns:
#   True if the task kind can create a model, False otherwise.
def kind_can_create_model(kind):
    if (
        kind == EnumKind.auto_ml.value
        or kind == EnumKind.finetune.value
        or kind == EnumKind.post_pretraining.value
        or kind == EnumKind.compress.value
        or kind == EnumKind.merge.value
        or kind == EnumKind.rlhf.value
    ):
        return True
    return False